import requests
from loguru import logger
from typing import Dict, List, Any, Optional
from config import DIFY_API_BASE, DIFY_DATASETS_ENDPOINT, DIFY_RETRIEVE_ENDPOINT, DIFY_API_KEY

class DifyKnowledgeBase:
    """Dify知识库工具类"""
    
    @staticmethod
    def get_datasets() -> List[Dict]:
        """
        获取所有可用的知识库列表
        :return: 知识库列表
        """
        try:
            url = f"{DIFY_API_BASE}{DIFY_DATASETS_ENDPOINT}"
            headers = {
                "Authorization": f"Bearer {DIFY_API_KEY}"
            }
            response = requests.get(url, headers=headers)
            response.raise_for_status()
            return response.json().get("data", [])
        except Exception as e:
            logger.error(f"获取知识库列表失败: {str(e)}")
            return []
    
    @staticmethod
    def retrieve_knowledge(dataset_id: str, query: str) -> Optional[Dict]:
        """
        从指定知识库检索相关内容
        :param dataset_id: 知识库ID
        :param query: 查询内容
        :return: 检索结果
        """
        try:
            url = f"{DIFY_API_BASE}{DIFY_RETRIEVE_ENDPOINT.format(dataset_id=dataset_id)}"
            headers = {
                "Authorization": f"Bearer {DIFY_API_KEY}",
                "Content-Type": "application/json"
            }
            payload = {
                "query": query,
                "retrieval_model": {
                    "search_method": "semantic_search",
                    "reranking_enable": False,
                    "reranking_model": {
                        "reranking_provider_name": "",
                        "reranking_model_name": ""
                    },
                    "top_k": 3,
                    "score_threshold_enabled": False,
                }
            }
            response = requests.post(url, headers=headers, json=payload)
            response.raise_for_status()
            return response.json()
        except Exception as e:
            logger.error(f"知识库检索失败: {str(e)}")
            return None
    
    @staticmethod
    def format_knowledge_result(result: Dict) -> str:
        """
        格式化知识库检索结果
        :param result: 检索结果
        :return: 格式化后的文本
        """
        if not result or "records" not in result:
            return ""
        
        formatted_text = "根据知识库内容：\n\n"
        for record in result["records"]:
            segment = record.get("segment", {})
            content = segment.get("content", "")
            if content:
                formatted_text += f"- {content}\n\n"
        
        return formatted_text 